#!/usr/bin/env python
from optparse import OptionParser
import os
import os.path
import sys

from matplotlib import pyplot

from charmicat import BadCommandLineError


def main():
    op = OptionParser(usage='%prog [opts] file.hpkl [file2.hpkl [...]]')

    op.add_option("-c", "--color", action="append", dest='colors', default=[],
                  help="Set first color with #<html color> (can be used N times to set first N colors)", 
                  metavar='<html color>')
    op.add_option("-e", "--ext", action='append', default=None, 
                  help="Draw output images with extension EXT (default='png', can be used more than once)",
                  metavar='EXT')
    op.add_option("-f", "--force", action="store_true", dest="force", default=False, 
                  help="Force overwrite of output directory")
    op.add_option("-n", "--name", action='append', dest='names', default=[],
                  help='Write images only for histogram NAME (can be used multiple times)',
                  metavar='NAME')
    op.add_option("-l", "--labels", dest='labels', default=None, 
                  help="Load Python dictionary in LABELS to provide labels for superimposed plots", 
                  metavar='LABELS')
    op.add_option("--no-legend", action='store_true', dest="no_legend", default=False,
                  help="Omit legend in final drawings")
    op.add_option("-o", "--output-dir", dest="output", default=None, 
                  help="Place output in directory DIR", metavar="DIR")    
    op.add_option("-p", "--presentation", action="store_true", dest="presentation",
                  default=False, help="Draw images in presentation mode")
    op.add_option("-t", "--tag", dest='tag', default='',
                  help='Append TAG to all histogram titles',
                  metavar='TAG')

    op.add_option("--alpha", type='float', dest='alpha', default=1.0, 
                  help='Set color alpha level to FRAC (b/t 0 and 1)', metavar="FRAC")


    opts, args = op.parse_args()

    mpl_colors = "#000000 #FF0000 #FF9966 #FFFF00 #009900 #66FF66 #0033FF #00CCFF #9933FF #9999FF #FF66FF #666666".split()

    if len(opts.colors) > 0:
        for i, c in enumerate(opts.colors):
            if c[0] == '#':
                mpl_colors[i] = c
            else:
                mpl_colors[i] = '#' + c


    if len(args) < 1:
        raise BadCommandLineError("Please supply at least one pickle file with histograms.", -1)
    elif len(args) > len(mpl_colors):
        msg = "Sorry, only {0} can be superimposed at a time ({1} requested)"
        raise BadCommandLineError(msg.format(len(mpl_colors), len(args)), -1)

    if opts.output is None:
        if args[0][-4:] == 'hpkl' or args[0][-4:] == '.pkl':
            opts.output = args[0].rsplit('.', 1)[0] + '.superimposed'
        else:
            opts.output = 'hists.superimposed'

    if os.access(opts.output, os.F_OK):
        if not os.path.isdir(opts.output):
            msg = "Output path `" + opts.output + "' already exists as something other than a directory.\n"
            msg += "Please specify another output path or remove the existing one."
            raise BadCommandLineError(msg, -2)

        elif not opts.force:
            msg  = "Output path `" + opts.output + "' already exists.\n"
            msg += "Please specify another output path or remove the existing one."
            raise BadCommandLineError(msg, -3)
    else:        
        os.mkdir(opts.output)

    if opts.ext is None:
        opts.ext = ['png']

    fontsize = 'medium'
    if opts.presentation:
        fontsize = 'xx-large'

    h = {}
    for inf in args:
        handle = os.path.basename(inf)

        if inf[-4:] == 'hpkl':
            handle = handle.rsplit('.', 1)[0]

        try:
            h[handle] = histogram.load(inf)
        except ValueError:
            print "Error in", inf
            raise

            
        fp.close()
        

    ignore = set()
    labels = {}
    
    attr = {}

    test_hdict = h[h.keys()[0]]

    hmod_load_and_apply(h)

    if '__hatt__' in test_hdict:
        try:
            analysis = __import__(test_hdict['__hatt__'][0])
            h_att_path = 'analysis.' + test_hdict['__hatt__'][1]
            h_att = eval(h_att_path)            
            

            for k, ih in h.iteritems():
                try:
                    labels[k] = ih['__label__']
                except KeyError:
                    labels[k] = k


                for key in ih:
                    for att in 'title xlabel xunits'.split():
                        if hasattr(ih[key], att) and key in h_att:
                            setattr(ih[key], att, getattr(h_att[key], att))
                    
                

                    # Remove 2D and 3D histograms from the draw list
                    if hasattr(ih[key], 'nybins'):
                        ignore.add(key)

            for k in test_hdict:
                attr[k] = {}
                for att in 'title xlabel xunits'.split():
                    if hasattr(test_hdict[k], att):
                        attr[k][att] = getattr(test_hdict[k], att)

            try:
                del test_hdict['__hatt__']
                del test_hdict['__label__']
            except KeyError:
                pass
            
        except ImportError, e:
            sys.stderr.write("Dictionary in `{0}' has `__hatt__' attribute but module `{1}'".format(args[0],
                                                                                                    test_hdict['__hatt__'][0]))
            sys.stderr.write(" specified therein cannot be imported:\n")
            sys.stderr.write(str(e))
            sys.stderr.write("Aborting...\n")
            return -4

    if len(opts.names):
        for name in opts.names:
            for k, ih in h.iteritems():
                if name not in ih:
                    msg  = "Histogram `{0}' requested with `-n' option but not found in `{1}'\n"
                    msg += "Aborting...\n"
                    sys.stderr.write(msg.format(name, args[i]))
                    return -5
            
        ih_keys = opts.names
    else:
        key_sets = [set([k2 for k2 in h[k] if not hasattr(h[k][k2], 'nybins') ]) for k in h]

        ih_keys = list(set.intersection(*key_sets))

    ih_keys.sort()
    h_keys = h.keys()

    test_hists = dict([(k, h[k][h[k].keys()[0]]) for k in h])

    h_keys.sort(key=lambda x: test_hists[x].nentries())
    h_keys.reverse()

    if opts.labels:
        with open(opts.labels) as fplbl:
            labels = eval(fplbl.read())

    
    legend_drawn=False
    legend_labels = []
    legend_patches = []
    for ikh, ih_key in enumerate(ih_keys):
        if opts.no_legend:
            fig = pyplot.figure(figsize=(8,6))    
            axes = pyplot.axes((0.125, 0.125, 0.8, 0.8))
        else:
            pyplot.figure(figsize=(14, 8.5))
            pyplot.axes((0.1, 0.075, 0.45, 0.75))
                
        for ik, hk in enumerate(h_keys):
            try:                
#                h[hk][ih_key].draw(name=os.path.join(opts.output, h[hk][ih_key].name), 
#                                   fontsize=fontsize, tag=opts.tag, 
#                                   hold=True, same=True, alpha='0.5', facecolor=mpl_colors[ik], 
#                                   label=labels[hk])

                label = "{0} ({1} events)".format(labels[hk], h[hk][ih_key].nentries())
                h[hk][ih_key].draw(name=os.path.join(opts.output, h[hk][ih_key].name), 
                                   fontsize=fontsize, tag=opts.tag, alpha=opts.alpha,
                                   hold=True, same=True, facecolor=mpl_colors[ik], 
                                   label=label)

                legend_labels.append(label)
                legend_patches.append(h[hk][ih_key].histret[-1][0])

            except ValueError, e:
                sys.stderr.write("Error drawing histogram `{0}': {1}\n".format(ih_key, e))
                sys.stderr.write("Skipping...\n")
                sys.stderr.flush()
                continue

#        if 'xunits' in attr[ih_key]:
#            print ih_key, attr[ih_key]['xlabel'], attr[ih_key]['xunits']


        if ih_key in attr:
            if 'xunits' in attr[ih_key] and attr[ih_key]['xunits'] is not None and attr[ih_key]['xunits'] not in attr[ih_key]['xlabel']:
                pyplot.xlabel("{0} ({1})".format(attr[ih_key]['xlabel'], attr[ih_key]['xunits']), fontsize=fontsize)
            else:
                pyplot.xlabel("{0}".format(attr[ih_key]['xlabel']), fontsize=fontsize)
        
            if 'title' in attr[ih_key] and attr[ih_key]['title'] is not None:
                pyplot.title(attr[ih_key]['title'])
            else:
                pyplot.title(attr[ih_key]['xlabel'])


        pyplot.ylabel(test_hdict[ih_key]._set_ylabel())


        if opts.no_legend:
            try:
                for ext in opts.ext:
                    pyplot.savefig(os.path.join(opts.output, h[hk][ih_key].name + '.' + ext))
            except RuntimeError, e:
                print "BLARG!", hk, ih_key, h[hk][ih_key].name
                raise 

            if not legend_drawn:
                pyplot.clf()
                pyplot.legend(legend_patches, legend_labels, loc='center')
                ax = pyplot.gca()
                ax.axis('off')
#                ax.axis('bottom').set_visible(False)

                for ext in opts.ext:
                    pyplot.savefig(os.path.join(opts.output, 'legend.' + ext))

                legend_drawn = True
                
        else:
            pyplot.legend(loc=(1.025, 0.075))
            

            try:
                for ext in opts.ext:
                    pyplot.savefig(os.path.join(opts.output, h[hk][ih_key].name + '.' + ext))
            except RuntimeError, e:
                print "BLARG!", hk, ih_key, h[hk][ih_key].name
                raise 
            
        pyplot.clf()

        sys.stdout.write('{0}/{1} complete\r'.format(ikh, len(ih_keys)))
        sys.stdout.flush()


    print '{0}/{1} complete'.format(len(ih_keys), len(ih_keys))



if __name__ == "__main__":
    try:
        sys.exit(main())
    except BadCommandLineError, e:
        sys.stderr.write(str(e) + "\n")
        sys.stderr.write("See `draw_superimposed.py -h' for more.\n")
        sys.stderr.flush()
        sys.exit(e.retcode)
